package cs470.kalman;

import java.awt.image.BufferedImage;
import java.util.ArrayList;

import Code.Point;
import com.panayotis.gnuplot.*;
import com.panayotis.gnuplot.JavaPlot.Key;
import com.panayotis.gnuplot.terminal.ImageTerminal;

import cs470.kalman.exception.KalmanException;

public class KalmanFilter 
{
	//TODO: verify deltaT value
	private static final double c = 0;		//linear friction
	private static final double deltaT = 0.5;	// time between computations

	private static final double bulletSpeed = 101;
	
	private static final double EPSILON = .000001;
	
	private Matrix H;
	private Matrix hTrans;
	private Matrix F;
	private Matrix fTrans;
	private Matrix sigmaX;	//TODO: fiddle with sigmaX matrix to improve performance
	private Matrix sigmaZ;
	
	//Matrices that change
	private Matrix K;
	private Matrix sigmaT;
	private Matrix mu;
	
	private double worldX;
	private double worldY;
	
	private KalmanPanel panel1;

	
	/************************************************
	 * Constructs a KalmanFilter object
	 * @param stdDev the stdDev of the noise
	 * @param z
	 * @param x the world width
	 * @param y the world height
	 * @throws Exception
	 */
	public KalmanFilter(double stdDev, Matrix z, double width, double height, final KalmanPanel panel1)
	{
		this.panel1 = panel1;

		worldX = width/2;
		worldY = height/2;
		
		try
		{
			H = new Matrix(new double [][]{
					{1, 0, 0, 0, 0, 0}, 
					{0, 0, 0, 1, 0, 0}});
			
			hTrans = H.transpose();
			
			double deltaTSquareHalved = deltaT*deltaT/2;
			F = new Matrix(new double [][]{
					{1, deltaT, deltaTSquareHalved, 0, 0, 0},
					{0, 1, deltaT, 0, 0, 0},
					{0, -1*c, 1, 0, 0, 0},
//					{0, 0, 1, 0, 0, 0},
					{0, 0, 0, 1, deltaT, deltaTSquareHalved},
					{0, 0, 0, 0, 1, deltaT},
					{0, 0, 0, 0, -1*c, 1}});
//					{0, 0, 0, 0, 0, 1}});

			
			fTrans = F.transpose();
			
			//Covariance
			sigmaX = new Matrix(new double [][]{
					{.1, 0, 0, 0, 0, 0},
					{0, .1, 0, 0, 0, 0},
					{0, 0, 1, 0, 0, 0},
					{0, 0, 0, .1, 0, 0},
					{0, 0, 0, 0, .1, 0},
					{0, 0, 0, 0, 0, 1}});
			sigmaZ = new Matrix(new double [][]{
					{stdDev*stdDev, 0},
					{0, stdDev*stdDev}});
			//TODO: look at tweaking sigma0 values
			sigmaT = new Matrix(new double[][]{
					{100, 0, 0, 0, 0, 0},
					{0, 0.1, 0, 0, 0, 0},
					{0, 0, 0.1, 0, 0, 0},
					{0, 0, 0, 100, 0, 0},
					{0, 0, 0, 0, 0.1, 0},
					{0, 0, 0, 0, 0, 0.1}});
			mu = new Matrix(new double[][]{{0},{0},{0},{0},{0},{0}});
			Matrix temp = calcFSigmaTFTrans().add(sigmaX);
			calcK(temp);
			calcMu(z);
			calcSigmaT(temp);
		}
		catch(Exception e)
		{
			System.out.println(e.getMessage());
		}
	}

	private void calcSigmaT(Matrix FSigmaTFTrans) throws KalmanException 
	{
		Matrix temp = K.multiply(H);	//KH
		int idDimension = temp.matrix.length;
		Matrix identityMatrix = generateIDMatrix(idDimension);
		
		temp = identityMatrix.subtract(temp);	// (I- KH)
		sigmaT = temp.multiply(FSigmaTFTrans); 	//(I - KH)(FsigmatFT + Sigmax)
	}

	private Matrix generateIDMatrix(int idDimension) throws KalmanException {
		double[][] identity = new double[idDimension][idDimension];
		for (int i = 0; i < identity.length; i++)
		{
			for (int j = 0; j < identity.length; j++)
			{
				if (i == j)
					identity[i][j] = 1;
				else
					identity[i][j] = 0;
			}
		}
		Matrix identityMatrix = new Matrix (identity);
		return identityMatrix;
	}

	private void calcMu(Matrix z) throws KalmanException 
	{
		Matrix temp = F.multiply(mu);	//Fmu_t 
		Matrix temp2 = H.multiply(F);	// HF
		temp2 = temp2.multiply(mu);		// HFmu_t
		temp2 = z.subtract(temp2);		//z - HFmu
		temp2 = K.multiply(temp2);		//K(z - HFmu)
		mu = temp.add(temp2);			//Fmu + K(z - HFmu)
	}

	private void calcK(Matrix fSigmaTfTrans) throws KalmanException {
		Matrix temp1 = fSigmaTfTrans.multiply(hTrans);	// (FSigmatF^T + Sigmax)H^T
		Matrix temp2 = H.multiply(fSigmaTfTrans); 	//H(FSigmatF^T + Sigmax)
		temp2 = temp2.multiply(hTrans);			  	//H(FSigmatF^T + Sigmax)H^T
		temp2 = temp2.add(sigmaZ);					  	//H(FSigmatF^T + Sigmax)H^T + Sigmaz
		temp2 = temp2.invert();						//(H(FSigmatF^T + Sigmax)H^T + Sigmaz)^-1
		
		K = temp1.multiply(temp2);
	}
	
	private Matrix calcFSigmaTFTrans() throws KalmanException
	{
		Matrix temp = F.multiply(sigmaT);	//FSigma_tF^T
		return temp.multiply(fTrans);
	}
	
	public Point update(Matrix z, Point tankLoc)
	{
		try
		{
			Matrix temp = calcFSigmaTFTrans().add(sigmaX);
			calcK(temp);
			calcMu(z);
			calcSigmaT(temp);
			Matrix expectedPos = H.multiply(mu);
			Point currentGuess = new Point(expectedPos.getValue(0, 0), expectedPos.getValue(1, 0));
			
//			System.out.println(sigmaT.toString());
			
			Double predictions = currentGuess.calcDistance(tankLoc)/bulletSpeed/deltaT;
			int numPredictions = predictions.intValue() + 1;
			double weight = numPredictions - predictions;
					
			Matrix predictMu = mu;
			ArrayList<Point> predictedPts = new ArrayList<Point>();
			for (int i = 0; i < numPredictions; i++)
			{
//				System.out.println("predict Mu: " + predictMu.toString());
				predictMu = F.multiply(predictMu);
				Matrix prediction = H.multiply(predictMu);
				predictedPts.add(new Point(prediction.getValue(0, 0), prediction.getValue(1, 0)));
			}
			
			Point last = predictedPts.get(predictedPts.size() - 1);
			Point predicted;
			if(predictedPts.size() == 1) {
				predicted = last;
			} else {
				Point penultimate = predictedPts.get(predictedPts.size() - 2);
				double x = last.getX()*(1-weight) + penultimate.getX()*weight;
				double y = last.getY()*(1-weight) + penultimate.getY()*weight;
				predicted = new Point (x, y);
			}
//			System.out.println("last " + last.toString());
//			System.out.println("pen " + penultimate.toString());
			
			
//			System.out.println("guess " + currentGuess.toString());
//			System.out.println("predicted " + x + " , " + y);
			
	        BufferedImage raster = getPlotImage(currentGuess, true);
	        BufferedImage second = getPlotImage(predicted, false);

	        panel1.setImage(second, raster);
	        
	        return predicted;
		}
		catch(KalmanException e)
		{
			System.out.println(e);
			return new Point(0, 0);
		}
	}

	private int count = 0;
	
	private BufferedImage getPlotImage(Point currentGuess, boolean isGray) {
		final JavaPlot p = new JavaPlot(true);
        final ImageTerminal imageTerminal = new ImageTerminal();
        p.setTerminal(imageTerminal);
//        if(isGray) {
//        	p.set("palette", "model RGB functions 1-gray,1-gray,1-gray");
//        } else {
//        	p.set("palette", "rgbformulae 3,0,0");
//        }
        p.set("xrange", "[-400.0: 400.0]");
        p.set("yrange", "[-400.0: 400.0]");
        p.setKey(Key.OFF);
        p.set("pm3d", "");
        p.set("view", "map");
        p.set("size", "square");
        p.set("isosamples", "100");
        p.set("colorbox", "size 0,0");
        
        String plotStr = getPlotString(currentGuess, sigmaT.getValue(0, 0), sigmaT.getValue(3, 3));
        if(isGray) {
        	System.out.println(plotStr);
        }
        
        p.addPlot(plotStr);
        
        
        p.plot();
        
        return imageTerminal.getImage();
	}
	
	
	private String getPlotString(Point mean, double plotSigmaX, double plotSigmaY) {
		double rho = 0d;
		String plotStrFormat = "1.0/ (2.0 * pi * %f * %f * sqrt(1 - %f**2) )  * exp(-1.0/2.0 * ((x - %f)**2 / %f**2 + (y - %f)**2 / %f**2 - 2.0*%f*(x - %f)* (y - %f) /(%f*%f) ) ) with pm3d";
		String plotStr = String.format(plotStrFormat, 
				plotSigmaX + EPSILON, plotSigmaY + EPSILON, 
				rho, mean.getX(), 
				plotSigmaX + EPSILON, mean.getY(), plotSigmaY + EPSILON, 
				rho, mean.getX(), mean.getY(), 
				plotSigmaX + EPSILON, plotSigmaY + EPSILON);
		return plotStr;
	}

	
}
